
//% color="#FF7F50" iconWidth=50 iconHeight=40
namespace handwriting_recognition{

    //% block="加载数据" blockType="tag"
    export function initmodel() {}
    //% block="初始化模块" blockType="command"
    export function init(parameter: any, block: any) {

        Generator.addImport(`import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
import torchvision
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from sklearn.metrics import confusion_matrix, classification_report
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from PIL import Image
import os
import time
from torchinfo import summary
from shared import OptimizedMLP,weights_init,evaluate_model
from shared import test_q1,load_model,preprocess_image,segment_digits,create_layer
from shared import process_and_normalize_digits,recognize_digits,process_image
import warnings
warnings.filterwarnings("ignore")
pic_path = None
`)   

    }
    
    //% block="加载MNIST数据集路径[PATH]，返回训练集数据[TRAIN]与测试集数据[TEST]" blockType="command"
    //% PATH.shadow="string" PATH.defl="data"
    //% TRAIN.shadow="normal" TRAIN.defl=train_loader
    //% TEST.shadow="normal" TEST.defl=test_loader
    export function readcap1(parameter: any, block: any) {
        let path=parameter.PATH.code;
        let train=parameter.TRAIN.code;
        let test=parameter.TEST.code;
        Generator.addCode(`
transform_train = transforms.Compose([
    transforms.RandomRotation(10), 
    transforms.ToTensor(), # 将图像转换为张量
    transforms.Normalize((0.5,), (0.5,)) 
])

transform_test = transforms.Compose([ 
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_dataset = datasets.MNIST(root=${path}, train=True, download=True, transform=transform_train) 
test_dataset = datasets.MNIST(root=${path}, train=False, download=True, transform=transform_test) 

${train} = DataLoader(train_dataset, batch_size=64, shuffle=True) 
${test} = DataLoader(test_dataset, batch_size=1000, shuffle=False)`)
    }

    //% block="获取文件夹路径[PATH]的所有图像，结果存入变量[PIC]" blockType="command"
    //% PATH.shadow="string" PATH.defl="picture"
    //% PIC.shadow="normal" PIC.defl="data"
    export function initmodel1(parameter: any, block: any) {
        let path=parameter.PATH.code;
        let PIC=parameter.PIC.code;
        
        Generator.addCode(`${PIC} = [
    f
    for f in os.listdir(${path})
    if f.lower().endswith((".png", ".jpg", ".jpeg", ".bmp"))
]
pic_path = ${path}
`)

    }
    //% block="图像处理" blockType="tag"
    export function pic() {}
    //% block="保存图像预处理的中间结果至文件夹路径[PATH]" blockType="command"
    //% PATH.shadow="string" PATH.defl="intermediate_results"
    export function save_process_image(parameter: any, block: any) {
        let path=parameter.PATH.code;
        Generator.addCode(`intermediate_results_dir = ${path}
results_dict = []
results_dicta = {}
`)
    }


    //% block="对图像[PIC]执行二值化处理，[IMG]保存处理后图像，结果存入变量[VAR]" blockType="command"
    //% PIC.shadow="normal" PIC.defl="imgfile"
    //% IMG.shadow="dropdown" IMG.options="IMG"
    //% VAR.shadow="normal" VAR.defl="thresh_result"
    export function preprocess_image(parameter: any, block: any) {
        let PIC=parameter.PIC.code;
        let IMG=parameter.IMG.code;
        let VAR=parameter.VAR.code;
        Generator.addCode(`
if pic_path != None:  
    image_path = os.path.join(pic_path, ${PIC})
    image_name = ${PIC}
    ${VAR} = preprocess_image(
        image_path,
        save_intermediate=${IMG},
        output_dir=intermediate_results_dir,
        image_name=image_name,
    )
else:
    image_name = ${PIC}
    ${VAR} = preprocess_image(
        ${PIC},
        save_intermediate=${IMG},
        output_dir=intermediate_results_dir,
        image_name=image_name,
    )
`)

    }
    //% block="对图像[PATH]执行分割处理，[IMG]保存处理后图像，结果存入变量[VAR]" blockType="command"
    //% PATH.shadow="normal" PATH.defl="thresh_result"
    //% IMG.shadow="dropdown" IMG.options="IMG"
    //% VAR.shadow="normal" VAR.defl="digit_imgs"
    export function segment_digits(parameter: any, block: any) {
        let path=parameter.PATH.code;
        let IMG=parameter.IMG.code;
        let VAR=parameter.VAR.code;
        Generator.addCode(`
${VAR}  = segment_digits(
    ${path},
    save_intermediate=${IMG},
    output_dir=intermediate_results_dir,
    image_name=image_name,
)
`)
    }
    //% block="对图像[PATH]执行规范化处理，[IMG]保存处理后图像，结果存入变量[VAR]" blockType="command"
    //% PATH.shadow="normal" PATH.defl="digit_imgs"
    //% IMG.shadow="dropdown" IMG.options="IMG"
    //% VAR.shadow="normal" VAR.defl="processed_digits"
    export function process_and_normalize_digits(parameter: any, block: any) {
        let path=parameter.PATH.code;
        let IMG=parameter.IMG.code;
        let VAR=parameter.VAR.code;
        Generator.addCode(`
${VAR}  = process_and_normalize_digits(
    ${path},
    save_intermediate=${IMG},
    output_dir=intermediate_results_dir,
    image_name=image_name,
)
`)
    }
    //% block="调用模型[MODEL]对图像[PATH]分类，结果存入文件[FILE]" blockType="command"
    //% MODEL.shadow="normal" MODEL.defl="model"
    //% PATH.shadow="normal" PATH.defl="processed_digits"
    //% FILE.shadow="string" FILE.defl="output.txt"
    export function recognize_digits(parameter: any, block: any) {
        let model=parameter.MODEL.code;
        let path=parameter.PATH.code;
        let file=parameter.FILE.code;
        Generator.addInit("results_dict", `
image_output_dir = ${file}
f = open(image_output_dir, "w", encoding="utf-8")`)
        Generator.addCode(`
results_valve  = recognize_digits(${model}, ${path}, device)

results_str = ''.join(map(str, results_valve))  # 使用空字符串作为分隔符

f.write(f"{image_name}: {results_str}\\n")
results_list = image_name + ": " + results_str

`)
    }
    //% block="图像分类结果" blockType="reporter"
    export function results_dict(parameter: any, block: any) {
        Generator.addCode(`results_list`)

    }

    //% block="模型参数" blockType="tag"
    export function initmodelA() {}
    //% block="展平输入数据" blockType="command"
    export function Flatten(parameter: any, block: any) {
        Generator.addCode(`flatten = nn.Flatten()
layers = []
`)
    }
//     //% block="创建第[NUM]层全连接层，输入维度[IN]，输出维度[OUT]" blockType="command"
//     //% NUM.shadow="normal" NUM.defl=1
//     //% IN.shadow="normal" IN.defl=28*28
//     //% OUT.shadow="normal" OUT.defl=1024
//     export function Linear(parameter: any, block: any) {
//         let num=parameter.NUM.code;
//         let in1=parameter.IN.code;
//         let out1=parameter.OUT.code;
//         Generator.addCode(`input_dim${num} = ${in1}
// output_dim${num} = ${out1}        `)
//     }
//     //% block="对第[NUM]层全连接层 [OUT]维的输出执行批量归一化处理" blockType="command"
//     //% NUM.shadow="normal" NUM.defl=1
//     //% OUT.shadow="normal" OUT.defl=1024
//     export function BatchNorm1d(parameter: any, block: any) {
//         let num=parameter.NUM.code;
//         let out1=parameter.OUT.code;
//         Generator.addCode(`LeakyReLU${num} = ${out1}`)
//     }

//     //% block="利用LeakyReLU激活函数处理第[NUM]层全连接层输出，斜率[DIG]" blockType="command"
//     //% NUM.shadow="normal" NUM.defl=1
//     //% DIG.shadow="normal" DIG.defl="0.01"
//     export function LeakyReLU(parameter: any, block: any) {
//         let num=parameter.NUM.code;
//         let dig=parameter.DIG.code;
//         Generator.addCode(`dropout_rate${num} = ${dig}`)
//     }
//     //% block="对第[NUM]层全连接层执行Dropout操作，比例[DIG]" blockType="command"
//     //% NUM.shadow="normal" NUM.defl=1
//     //% DIG.shadow="normal" DIG.defl="0.1"
//     export function Dropout(parameter: any, block: any) {
//         let num=parameter.NUM.code;
//         let dig=parameter.DIG.code;
//         Generator.addCode(`dropout${num} = ${dig}
// layer_value${num}= create_layer(input_dim${num}, output_dim${num}, LeakyReLU${num}, dropout_rate${num})
// layers.extend(layer_value${num})`)
//     }


//         //% block="---" blockType="tag"
//         export function initmodelaBV() {}
    
    //% block="创建第[NUM]层全连接层，输入维度[IN]输出维度[OUT]" blockType="command"
    //% NUM.shadow="normal" NUM.defl="1"
    //% IN.shadow="normal" IN.defl=28*28
    //% OUT.shadow="normal" OUT.defl=1024
    export function model(parameter: any, block: any) {
        let NUM=parameter.NUM.code;
        let IN=parameter.IN.code;
        let OUT=parameter.OUT.code;
        if (NUM==1){
            Generator.addCode(`flatten = nn.Flatten()
layers = []`)
        }
        Generator.addCode(`
input_dim${NUM} = ${IN}
output_dim${NUM} = ${OUT}`)

    }
    

   //% block="第[NUM]层[OUT]维的输出执行批量归一化；执行LeakyReLU激活函数（斜率[DIG]）激活；执行Dropout操作（比例[DIGB]）" blockType="command"
    //% NUM.shadow="normal" NUM.defl="1"
    //% OUT.shadow="normal" OUT.defl=1024
    //% DIG.shadow="normal" DIG.defl="0.01"
    //% DIGB.shadow="normal" DIGB.defl="0.1"
    export function model1(parameter: any, block: any) {
        let NUM=parameter.NUM.code;
        let OUT=parameter.OUT.code;
        let DIG=parameter.DIG.code;
        let DIGB=parameter.DIGB.code;
        Generator.addCode(`output_dim${NUM} = ${OUT}
LeakyReLU${NUM} = ${DIG}
dropout_rate${NUM} = ${DIGB}
layer_value${NUM}= create_layer(input_dim${NUM}, output_dim${NUM}, LeakyReLU${NUM}, dropout_rate${NUM})
layers.extend(layer_value${NUM})
`)

    }



        //% block="---" blockType="tag"
        export function initmodelaB() {}
    //% block="创建输出层，输入维度[IN]，输出维度[OUT]" blockType="command"
    //% IN.shadow="normal" IN.defl=64
    //% OUT.shadow="normal" OUT.defl=10
    export function Linear1(parameter: any, block: any) {
        let in1=parameter.IN.code;
        let out1=parameter.OUT.code;
        Generator.addCode(`fc6 = nn.Linear(${in1}, ${out1})
layers.append(fc6)        
Optimized = OptimizedMLP(flatten,layers)`)
    }
        //% block="模型训练" blockType="tag"
        export function initmodela() {}
    //% block="定义损失函数[A]" blockType="command"
    //% A.shadow="normal" A.defl="criterion"
    export function CrossEntropyLoss(parameter: any, block: any) {
        let path=parameter.A.code;
        Generator.addCode(`${path} = nn.CrossEntropyLoss()\n`) 
    
    }
    //% block="定义优化器[A]" blockType="command"
    //% A.shadow="normal" A.defl="optimizer"
    export function optimizer(parameter: any, block: any) {
        let path=parameter.A.code;
        Generator.addCode(`model_opt = Optimized.to(device) 
model_opt.apply(weights_init)
${path} = optim.Adam(model_opt.parameters(), lr=0.001, weight_decay=1e-5) 
scheduler = StepLR(${path}, step_size=5, gamma=0.5)
`) 
    
    }
    
    //% block="用损失函数[A]与优化器[B]，对训练集数据[TRAIN]与测试集数据[TEST]进行多层感知器模型训练" blockType="command"
    //% A.shadow="normal" A.defl=criterion
    //% B.shadow="normal" B.defl=optimizer
    //% TRAIN.shadow="normal" TRAIN.defl=train_loader
    //% TEST.shadow="normal" TEST.defl=test_loader
    
    export function CrossEntropyLoss1(parameter: any, block: any) {
        let a=parameter.A.code;
        let b=parameter.B.code;
        let train=parameter.TRAIN.code;
        let test=parameter.TEST.code;

        Generator.addCode(`

for epoch in range(num_epochs):
    model_opt.train() 
    running_loss = 0.0 

    for images, labels in ${train}: 
        images = images.view(-1, 28*28).to(device)  
        labels = labels.to(device)  

        ${b}.zero_grad()   
        outputs = model_opt(images) 
        loss = ${a}(outputs, labels)      
        loss.backward()        
        ${b}.step()     

        running_loss += loss.item()     
    accuracy = evaluate_model(model_opt, ${test}, device)
    scheduler.step()
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}, Accuracy: {accuracy:.2f}%') # 打印每个 epoch 的损失值和准确率。 
`)

    }
    //% block="训练轮数[NUM]，选择[MODEL]算力" blockType="command"
    //% MODEL.shadow="dropdown" MODEL.options="MODEL"
    //% NUM.shadow="normal" NUM.defl=5
    
    export function num_epochs(parameter: any, block: any) {

        let model=parameter.MODEL.code;
        let num=parameter.NUM.code;
        Generator.addInit("device", `device = torch.device('${model}')
num_epochs = ${num}  `)

    }
    //% block="---" blockType="tag"
    export function initmodelaBQ() {}
    //% block="打印模型的结构信息" blockType="command"
    export function initmodel3(parameter: any, block: any) {
        Generator.addCode(`test_q1(Optimized)`)
    }
    //% block="模型在测试集上[TEST]的准确率" blockType="reporter"
        //% TEST.shadow="normal" TEST.defl=test_loader
    export function accuracy(parameter: any, block: any) {
        let test=parameter.TEST.code;
        Generator.addCode(`evaluate_model(model_opt, test_loader, device)`)


}     


    //% block="---" blockType="tag"
    export function initmodelaBQQ() {}
    //% block="保存模型路径[PATH]" blockType="command"
    //% PATH.shadow="string" PATH.defl="optimized_mlp_mnist.pth"
    export function save(parameter: any, block: any) {
        let path=parameter.PATH.code;
        Generator.addCode(`torch.save(model_opt, ${path})`) 
    
    }
    //% block="加载模型路径[PATH]选择[MODEL]算力，结果存入变量[VAR]" blockType="command"
    //% PATH.shadow="string" PATH.defl="optimized_mlp_mnist.pth"
    //% MODEL.shadow="dropdown" MODEL.options="MODEL"
    //% VAR.shadow="normal" VAR.defl="model"
    export function load_model(parameter: any, block: any) {
        let path=parameter.PATH.code;
        let model=parameter.MODEL.code;
        let var1=parameter.VAR.code;
        Generator.addCode(`
device = torch.device('${model}')
${var1} = load_model(${path}, device)`) 
    
    }


}